#!/usr/bin/env python3
'''
未验证

目前测试发现是可以缓慢提高的，但是不知道最终能到多高，可能需要非常长的训练时间
训练到一般发现奖励不在上升，添加了学习率衰减器，奖励进一步上升，可能是由于学习率迭代器迭代太快，导致奖励上升缓慢
减缓迭代器的迭代
在添加迭代器前，对模型进行了保存，后续如果迭代器不对可以讲模型进行恢复，重新进行训练
完成代码调整，准备重新训练
训练记录：
20250213: 训练分数81.40，测试分数83.10分，发现保存最好的代码以及训练中的代码存在问题，需要调整代码后继续训练
20250214: Starting from step 54000,当前学习率： 0.00045000000000000004

训练分数94.6，测试分数94.6，继续训练

20250217： 0.00040500000000000003，训练分数98.9，测试分数98.9，继续训练
20250218:当前学习率： 0.00032805000000000003，训练、测试分数105.10，继续训练
20250219:当前学习率： 0.000295245,训练分数、测试分数104.8，继续训练一天，关闭调度器，如果无法提高，那么就重头训练，不开启调度器
20250220:当前学习率： 0.00023914845，训练测试分数109.4，继续训练
20250221:当前学习率： 0.00023914845,训练测试分数122.6，继续训练
20250222:当前学习率： 0.00023914845，测试训练分数174.2，继续训练
20250224:当前学习率： 0.00023914845,训练测试分数185.10，继续训练，并解决模型（最好模型和训练模型只能保存一份）保存问题
20250225:当前学习率： 0.00023914845,训练测试分数207.9，继续训练
20250226:当前学习率： 0.00023914845，训练测试分数263.60，继续训练
20250227:当前学习率： 0.00023914845，训练测试分数317.40，继续训练
20250301:当前学习率： 0.00023914845，训练测试分数327.50，继续训练
20250302:当前学习率： 0.00023914845，训练测试分数402.40，继续训练，增加保存训练过程中找到的最好模型
20250303:当前学习率： 0.00023914845,训练测试分数196.3分，继续训练，开放学习率调度器
20250304：当前学习率： 0.00023914845，由于训练终端，测试训练分数未知，继续训练
20250305:当前学习率： 0.000215233605,训练测试分数149。6分，继续训练一天
20250306：当前学习率： 0.00017433922005,训练测试分数166.6分，继续训练一天，无进步则停止训练
20250307：当前学习率： 0.00014121476824,训练测试分数168.5分，关闭学习率调度器，继续训练
20250308:当前学习率： 0.00011438396227480502，训练测试分数181分，关闭学习率调度器，继续训练
20250309:当前学习率： 0.0001143839622748050，训练测试分数103.80分，开放学习率调度器，继续训练
200250310：当前学习率： 0.00011438396227480502,训练测试分数121分，停止训练，调增代码或者重头训练
'''
import os
import ptan
import time
import argparse
from tensorboardX import SummaryWriter
import gymnasium as gym
import torch
import torch.optim as optim

from lib import common
import ale_py

gym.register_envs(ale_py)


LEARNING_RATE = 5e-4
TEST_EVERY_BATCH = 1000
SAVE_EVERY_BATCH = 100

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--cuda", default=True, action="store_true", help="Enable cuda")
    parser.add_argument("-n", "--name", default='alien', help="Name of the run")
    parser.add_argument("--steps", type=int, default=None, help="Limit of training steps, default=disabled")
    args = parser.parse_args()
    device = common.select_device(args=args)

    saves_path = os.path.join("saves", "01_a2c_" + args.name)
    os.makedirs(saves_path, exist_ok=True)

    envs = [common.make_env() for _ in range(common.NUM_ENVS)]

    test_env = common.make_env()
    writer = SummaryWriter(comment="-01_a2c_" + args.name)

    # 创建网络
    net = common.AtariA2C(envs[0].observation_space.shape, envs[0].action_space.n).to(device)
    print(net)

    # todo 这里为什么采用了RMSprop优化器
    optimizer = optim.RMSprop(net.parameters(), lr=LEARNING_RATE, eps=1e-5)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=50000, gamma=0.9)

    step_idx = 0
    total_steps = 0
    best_reward = None
    best_test_reward = None

    # checkpoint_path = os.path.join(saves_path, "checkpoint.pth")
    # Load checkpoint if exists
    # if os.path.exists(checkpoint_path):
    if os.path.exists(saves_path) and len(os.listdir(saves_path)) > 0:
         # 增加加载模型的代码
        checkpoints = sorted(filter(lambda x: "epoch" in x, os.listdir(saves_path)),
                             key=lambda x: int(x.split('_')[-1].split('.')[0]))
        if len(checkpoints) > 0:
            checkpoint = torch.load(os.path.join(saves_path, checkpoints[-1]), map_location=device, weights_only=False)
            net.load_state_dict(checkpoint['model_state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
            step_idx = checkpoint['step_idx']
            total_steps = checkpoint['total_steps']
            scheduler.load_state_dict(checkpoint['scheduler_state_dict'])

            print(f"加载模型成功. Starting from step {step_idx}")
            print(f"当前学习率：", scheduler.get_last_lr()[0])

    ts_start = time.time()
    with ptan.common.utils.TBMeanTracker(writer, batch_size=100) as tb_tracker:
        # 遍历所有的游戏环境，使用现有的net，预测出REWARD_STEPS的动作以及动作对应的奖励、状态价值
        for mb_obs, mb_rewards, mb_actions, mb_values, _, done_rewards, done_steps in \
                common.iterate_batches(envs, net, device=device):
            if len(done_rewards) > 0:
                # 如果有收集到数据，那么计算公共执行的步数
                total_steps += sum(done_steps)
                # 计算执行步数的速度
                speed = total_steps / (time.time() - ts_start)
                # 保存记录最好的回报奖励
                if best_reward is None:
                    best_reward = done_rewards.max()
                elif best_reward < done_rewards.max():
                    best_reward = done_rewards.max()
                    common.save_best_model(best_reward, net.state_dict(), saves_path, f"best_{step_idx}_")
                # 记录到tensorboard中
                tb_tracker.track("total_reward_max", best_reward, step_idx)
                tb_tracker.track("total_reward", done_rewards, step_idx)
                tb_tracker.track("total_steps", done_steps, step_idx)
                print("%d: done %d episodes, mean_reward=%.2f, best_reward=%.2f, speed=%.2f" % (
                    step_idx, len(done_rewards), done_rewards.mean(), best_reward, speed))

            # 训练网络
            common.train_a2c(net, mb_obs, mb_rewards, mb_actions, mb_values,
                             optimizer, tb_tracker, step_idx, device=device)
            step_idx += 1
            scheduler.step()
            # 判断步数结束标识
            if args.steps is not None and args.steps < step_idx:
                break

            if step_idx % TEST_EVERY_BATCH == 0:
                # 测试网络
                net.eval()
                test_reward, test_steps = common.test_model(test_env, net, device=device)
                net.train()
                # 记录测试结果
                writer.add_scalar("test_reward", test_reward, step_idx)
                writer.add_scalar("test_steps", test_steps, step_idx)
                if best_test_reward is None or best_test_reward < test_reward:
                    # 根据测试结果，保存最好的模型
                    best_test_reward = test_reward
                common.save_best_model(test_reward, net.state_dict(), saves_path, f"best_{step_idx}_")
                print("%d: test reward=%.2f, steps=%.2f, best_reward=%.2f" % (
                    step_idx, test_reward, test_steps, best_test_reward))

            if step_idx % SAVE_EVERY_BATCH == 0:
                # Save checkpoint
                checkpoint = {
                    'step_idx': step_idx,
                    'model_state_dict': net.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                    'total_steps': total_steps,
                    'scheduler_state_dict': scheduler.state_dict()
                }

                common.save_checkpoints(step_idx, checkpoint, saves_path, "01_a2c")
                    # torch.save({
                    #     'step_idx': step_idx,
                    #     'model_state_dict': net.state_dict(),
                    #     'optimizer_state_dict': optimizer.state_dict(),
                    #     'total_steps': total_steps,
                    #     'scheduler_state_dict': scheduler.state_dict()
                    # }, checkpoint_path)
